#include <sys/reboot.h>
#include "upgrade.h"
#include <curl/curl.h>
#include <fstream>
#include <iostream>
#include <thread>
#include "common_config_manager.h"
#include "base_log.h"
#include "base_utils.h"
#include <algorithm>
#include <openssl/sha.h>
#include <openssl/evp.h>
#include <sys/mount.h>
#include <sys/stat.h>
#include <unistd.h>
#include <iomanip>
#include <minizip-ng/mz.h>
#include <minizip-ng/mz_zip.h>
#include <minizip-ng/mz_strm.h>
#include <minizip-ng/mz_strm_mem.h>
#include <minizip-ng/mz_zip_rw.h>
#include <minizip-ng/mz_strm_os.h>
#include <sys/ioctl.h>
#include <mtd/mtd-user.h>
#include <fcntl.h>
#include "flash_manager.h"
#include "upgrade_diagnosis.h"

namespace {
bool DropSystemCaches()
{
    // 同步文件系统，确保所有挂起的写操作完成
    sync();

    // 写入3会清除 pagecache、dentries 和 inodes 缓存
    std::ofstream drop_caches("/proc/sys/vm/drop_caches");
    if (!drop_caches) {
        LOG_ERROR("Failed to open /proc/sys/vm/drop_caches");
        return false;
    }

    if (!(drop_caches << "3" << std::flush)) {
        LOG_ERROR("Failed to write to /proc/sys/vm/drop_caches");
        return false;
    }

    return true;
}
} // namespace

namespace El {
namespace Upgrade {

Upgrade::Upgrade() : status_(UpgradeStatus::IDLE), progress_(0) {}

Upgrade::~Upgrade()
{
    // 空实现
}

bool Upgrade::Start()
{
    // 注册升级命令
    UpgradeDiagnosis::GetInstance().Start();
    return true;
}

void Upgrade::Stop()
{
    // 空实现
}

Upgrade &Upgrade::GetInstance()
{
    static Upgrade instance;
    return instance;
}

// 设置进度回调函数
void Upgrade::SetProgressCallback(std::function<void(int, bool)> callback)
{
    progressCallback_ = callback;
}

bool Upgrade::PrepareUpgrade()
{
    // 释放系统缓存
    if (!DropSystemCaches()) {
        LOG_WARN("Failed to drop system caches, continuing anyway");
    }

    // 清空并预分配buffer
    buffer_.clear();
    try {
        buffer_.reserve(10 * 1024 * 1024);
    } catch (const std::bad_alloc &e) {
        LOG_ERROR("Failed to reserve buffer memory");
        return false;
    }

    return true;
}

bool Upgrade::Append(const uint8_t *data, size_t length)
{
    if (data == nullptr || length == 0) {
        LOG_ERROR("Failed to append data");
        return false;
    }

    try {
        // 追加数据到缓冲区
        size_t oldSize = buffer_.size();
        buffer_.resize(oldSize + length);
        std::memcpy(buffer_.data() + oldSize, data, length);
        return true;
    } catch (const std::exception &e) {
        LOG_ERROR("Failed to append data");
        return false;
    }
}

bool Upgrade::StartUpgrade()
{
    if (buffer_.empty()) {
        LOG_ERROR("No upgrade package data");
        status_ = UpgradeStatus::FAILED;
        if (progressCallback_) {
            progressCallback_(progress_, false); // 失败
        }
        return false;
    }

    // 开始解析和验证升级包
    status_ = UpgradeStatus::VERIFYING;
    progress_ = 0;
    if (progressCallback_) {
        progressCallback_(progress_, true); // 成功
    }

    if (!ParseUpgradePackage() || !VerifyUpgradePackage()) {
        LOG_ERROR("Package verification failed");
        status_ = UpgradeStatus::FAILED;
        if (progressCallback_) {
            progressCallback_(progress_, false); // 失败
        }
        return false;
    }

    progress_ = 5;
    if (progressCallback_) {
        progressCallback_(progress_, true);
    }

    status_ = UpgradeStatus::UPGRADING;

    // 执行pre脚本（不计入进度）
    if (packageInfo_.scripts.find("pre") != packageInfo_.scripts.end()) {
        if (!ExecuteScript(packageInfo_.scripts["pre"])) {
            LOG_ERROR("Failed to execute pre script");
            status_ = UpgradeStatus::FAILED;
            if (progressCallback_) {
                progressCallback_(progress_, false);
            }
            return false;
        }
    }

    // 计算所有需要更新的文件的总大小，用于进度权重计算
    size_t totalSize = 0;
    for (const auto &fileInfo : packageInfo_.fileInfos) {
        if (fileInfo.need_update) {
            totalSize += fileInfo.size;
        }
    }

    float currentProgress = 5.0f;          // 从5%开始
    const float totalFileProgress = 95.0f; // 文件操作总共占95%

    // 写入每个文件
    for (const auto &fileInfo : packageInfo_.fileInfos) {
        if (!fileInfo.need_update) {
            LOG_INFO("Skipping file {}: update not needed", fileInfo.name);
            continue;
        }

        // 根据文件大小计算该文件占总进度的比例
        float fileWeight = totalSize > 0 ? (static_cast<float>(fileInfo.size) / totalSize) : 0;
        float fileProgressRange = fileWeight * totalFileProgress;

        // 文件写入占70%，验证占30%
        float writeRange = fileProgressRange * 0.7f;

        // 设置写入进度回调
        auto writeProgressCallback = [&](int percent) {
            progress_ = static_cast<int>(currentProgress + (writeRange * percent / 100));
            if (progressCallback_) {
                progressCallback_(progress_, true);
            }
        };

        void *handle = mz_zip_reader_create();
        if (!handle) {
            LOG_ERROR("Failed to create zip reader");
            status_ = UpgradeStatus::FAILED;
            if (progressCallback_) {
                progressCallback_(progress_, false);
            }
            return false;
        }

        if (mz_zip_reader_open_buffer(handle, buffer_.data(), buffer_.size(), 0) != MZ_OK) {
            LOG_ERROR("Failed to open zip buffer");
            mz_zip_reader_delete(&handle);
            status_ = UpgradeStatus::FAILED;
            if (progressCallback_) {
                progressCallback_(progress_, false);
            }
            return false;
        }

        if (mz_zip_reader_locate_entry(handle, fileInfo.name.c_str(), 0) != MZ_OK) {
            LOG_ERROR("Failed to locate entry: {}", fileInfo.name);
            mz_zip_reader_delete(&handle);
            status_ = UpgradeStatus::FAILED;
            if (progressCallback_) {
                progressCallback_(progress_, false);
            }
            return false;
        }

        if (mz_zip_reader_entry_open(handle) != MZ_OK) {
            LOG_ERROR("Failed to open entry: {}", fileInfo.name);
            mz_zip_reader_delete(&handle);
            status_ = UpgradeStatus::FAILED;
            if (progressCallback_) {
                progressCallback_(progress_, false);
            }
            return false;
        }

        // 使用FlashManager
        El::Upgrade::FlashManager manager;
        if (!manager.Open(fileInfo.partition)) {
            LOG_ERROR("Failed to open partition: {}", fileInfo.partition);
            mz_zip_reader_delete(&handle);
            status_ = UpgradeStatus::FAILED;
            if (progressCallback_) {
                progressCallback_(progress_, false);
            }
            return false;
        }

        auto info = manager.GetInfo();
        std::vector<char> buf(info.eraseSize);
        size_t offset = 0;

        // 先验证当前flash中的数据是否与要写入的数据相同
        bool needWrite = false;
        void *verifyHandle = mz_zip_reader_create();
        if (!verifyHandle) {
            LOG_ERROR("Failed to create zip reader for verification");
            manager.Close();
            mz_zip_reader_delete(&handle);
            status_ = UpgradeStatus::FAILED;
            if (progressCallback_) {
                progressCallback_(progress_, false);
            }
            return false;
        }

        if (mz_zip_reader_open_buffer(verifyHandle, buffer_.data(), buffer_.size(), 0) != MZ_OK ||
            mz_zip_reader_locate_entry(verifyHandle, fileInfo.name.c_str(), 0) != MZ_OK ||
            mz_zip_reader_entry_open(verifyHandle) != MZ_OK) {
            LOG_ERROR("Failed to prepare zip file for verification");
            mz_zip_reader_delete(&verifyHandle);
            manager.Close();
            mz_zip_reader_delete(&handle);
            status_ = UpgradeStatus::FAILED;
            if (progressCallback_) {
                progressCallback_(progress_, false);
            }
            return false;
        }

        size_t verifiedSize = 0;
        std::vector<char> flashBuf(info.eraseSize);
        std::vector<char> zipBuf(info.eraseSize);

        while (verifiedSize < fileInfo.size) {
            size_t toRead = std::min(info.eraseSize, fileInfo.size - verifiedSize);

            // 计算当前文件的验证进度百分比
            float fileVerifyPercent = static_cast<float>(verifiedSize) / fileInfo.size;
            // 验证阶段占整个文件进度的30%
            progress_ = static_cast<int>(currentProgress + (fileProgressRange * 0.3f * fileVerifyPercent));
            if (progressCallback_) {
                progressCallback_(progress_, true);
            }

            // 读取flash中的数据
            ssize_t flashReadCount = manager.Read(reinterpret_cast<uint8_t *>(flashBuf.data()), toRead, verifiedSize);
            if (flashReadCount <= 0) {
                LOG_ERROR("Failed to read flash data for comparison");
                mz_zip_reader_delete(&verifyHandle);
                manager.Close();
                mz_zip_reader_delete(&handle);
                status_ = UpgradeStatus::FAILED;
                if (progressCallback_) {
                    progressCallback_(progress_, false);
                }
                return false;
            }

            // 读取zip文件中的数据
            int32_t zipReadCount = mz_zip_reader_entry_read(verifyHandle, zipBuf.data(), toRead);
            if (zipReadCount <= 0 || zipReadCount != flashReadCount) {
                LOG_ERROR("Failed to read zip data for comparison or size mismatch");
                mz_zip_reader_delete(&verifyHandle);
                manager.Close();
                mz_zip_reader_delete(&handle);
                status_ = UpgradeStatus::FAILED;
                if (progressCallback_) {
                    progressCallback_(progress_, false);
                }
                return false;
            }

            // 比较数据
            if (memcmp(flashBuf.data(), zipBuf.data(), flashReadCount) != 0) {
                needWrite = true;
                break;
            }

            verifiedSize += flashReadCount;
        }

        mz_zip_reader_delete(&verifyHandle);

        if (!needWrite) {
            LOG_INFO("File {} content is identical, skipping write", fileInfo.name);
            manager.Close();
            mz_zip_reader_delete(&handle);

            // 更新当前总进度
            currentProgress += fileProgressRange;
            progress_ = static_cast<int>(currentProgress);
            if (progressCallback_) {
                progressCallback_(progress_, true);
            }
            continue;
        }

        LOG_INFO("File {} content is different, performing write", fileInfo.name);

        // 重新打开文件准备写入
        if (mz_zip_reader_entry_close(handle) != MZ_OK || mz_zip_reader_entry_open(handle) != MZ_OK) {
            LOG_ERROR("Failed to seek to start of file for writing");
            manager.Close();
            mz_zip_reader_delete(&handle);
            status_ = UpgradeStatus::FAILED;
            if (progressCallback_) {
                progressCallback_(progress_, false);
            }
            return false;
        }

        // 写入数据
        int32_t bytesRead;
        size_t totalWritten = 0;
        offset = 0;

        while ((bytesRead = mz_zip_reader_entry_read(handle, buf.data(), info.eraseSize)) > 0) {
            if (!manager.Write(reinterpret_cast<const uint8_t *>(buf.data()), bytesRead, offset,
                               writeProgressCallback)) {
                LOG_ERROR("Write failed at offset {}", offset);
                manager.Close();
                mz_zip_reader_delete(&handle);
                status_ = UpgradeStatus::FAILED;
                if (progressCallback_) {
                    progressCallback_(progress_, false);
                }
                return false;
            }
            offset += bytesRead;
            totalWritten += bytesRead;
            if (totalWritten > fileInfo.size) {
                LOG_ERROR("File size exceeds declared size: {}", fileInfo.name);
                manager.Close();
                mz_zip_reader_delete(&handle);
                status_ = UpgradeStatus::FAILED;
                if (progressCallback_) {
                    progressCallback_(progress_, false);
                }
                return false;
            }
        }

        manager.Close();
        mz_zip_reader_delete(&handle);

        LOG_INFO("Upgrade {} Success", fileInfo.name);

        // 更新当前总进度
        currentProgress += fileProgressRange;
        progress_ = static_cast<int>(currentProgress);
        if (progressCallback_) {
            progressCallback_(progress_, true);
        }
    }

    // 执行post脚本（不计入进度）
    if (packageInfo_.scripts.find("post") != packageInfo_.scripts.end()) {
        if (!ExecuteScript(packageInfo_.scripts["post"])) {
            LOG_ERROR("Failed to execute post script");
            status_ = UpgradeStatus::FAILED;
            if (progressCallback_) {
                progressCallback_(progress_, false);
            }
            return false;
        }
    }

    status_ = UpgradeStatus::SUCCESS;
    progress_ = 100;
    if (progressCallback_) {
        progressCallback_(progress_, true); // 成功
    }

    // 立即重启
    reboot(RB_AUTOBOOT);
    return true;
}

void Upgrade::StartUpgradeAsync()
{
    // 创建一个线程来执行升级过程，不等待结果
    std::thread([this]() { StartUpgrade(); }).detach();
}

bool Upgrade::ParseUpgradePackage()
{
    void *handle = mz_zip_reader_create();
    if (handle == nullptr) {
        LOG_ERROR("Failed to create zip reader");
        return false;
    }

    if (mz_zip_reader_open_buffer(handle, buffer_.data(), buffer_.size(), 0) != MZ_OK) {
        LOG_ERROR("Failed to open zip buffer");
        mz_zip_reader_delete(&handle);
        return false;
    }

    // 读取manifest.json
    if (mz_zip_reader_locate_entry(handle, "manifest.json", 0) != MZ_OK) {
        mz_zip_reader_delete(&handle);
        throw std::runtime_error("manifest.json not found in package");
    }

    if (mz_zip_reader_entry_open(handle) != MZ_OK) {
        mz_zip_reader_delete(&handle);
        throw std::runtime_error("Failed to open manifest.json");
    }

    std::string manifest_content;
    char buf[8192];
    int32_t bytesRead;

    while ((bytesRead = mz_zip_reader_entry_read(handle, buf, sizeof(buf))) > 0) {
        manifest_content.append(buf, bytesRead);
    }

    mz_zip_reader_entry_close(handle);
    mz_zip_reader_delete(&handle);

    if (!ParseManifest(manifest_content)) {
        LOG_ERROR("Failed to parse manifest.json");
        return false;
    }

    return true;
}

bool Upgrade::ParseManifest(const std::string &manifestContent)
{
    try {
        auto json = nlohmann::json::parse(manifestContent);

        packageInfo_.productName = json["product"]["name"];
        packageInfo_.hardwareVersion = json["product"]["hardware"];
        packageInfo_.toVersion = json["version"]["to"];

        if (json["version"].contains("from")) {
            packageInfo_.fromVersion = json["version"]["from"];
            packageInfo_.type = UpgradeType::DIFF;
        } else {
            packageInfo_.type = UpgradeType::FULL;
        }

        packageInfo_.minVersion = json["compatibility"]["min_version"];
        packageInfo_.maxVersion = json["compatibility"]["max_version"];

        // 解析文件列表
        packageInfo_.fileInfos.clear();
        for (const auto &file : json["files"]) {
            FileUpgradeInfo fileInfo;
            fileInfo.name = file["name"];
            fileInfo.partition = file["partition"];
            fileInfo.offset = file["offset"];
            fileInfo.size = file["size"];
            fileInfo.verify = file["verify"];
            fileInfo.hash = file["hash"];
            // 添加need_update字段解析,默认为true
            fileInfo.need_update = file.value("need_update", true);

            LOG_INFO("File {} size: {}", fileInfo.name.c_str(), fileInfo.size);

            // 移除hash值中的"sha256:"前缀
            if (fileInfo.hash.substr(0, 7) == "sha256:") {
                fileInfo.hash = fileInfo.hash.substr(7);
            }

            packageInfo_.fileInfos.push_back(fileInfo);
        }

        // 解析脚本配置
        if (json.contains("scripts") && json["scripts"]["enabled"]) {
            if (json["scripts"].contains("pre")) {
                packageInfo_.scripts["pre"] = json["scripts"]["pre"];
            }
            if (json["scripts"].contains("post")) {
                packageInfo_.scripts["post"] = json["scripts"]["post"];
            }
            if (json["scripts"].contains("rollback")) {
                packageInfo_.scripts["rollback"] = json["scripts"]["rollback"];
            }
        }

        return true;
    } catch (const std::exception &e) {
        LOG_ERROR("Failed to parse manifest: {}", e.what());
        return false;
    }
}

bool Upgrade::VerifyUpgradePackage()
{
    return true;

    // 验证产品型号
    auto productConfig = El::Common::ConfigManager::GetInstance().GetConfig("product");
    if (productConfig.is_null() || !productConfig.contains("name")) {
        LOG_ERROR("Product name not configured");
        return false;
    }
    std::string productName = productConfig["name"].get<std::string>();
    if (productName != packageInfo_.productName) {
        LOG_ERROR("Product name mismatch");
        return false;
    }

    // 验证硬件版本
    if (!productConfig.contains("hardware_version")) {
        LOG_ERROR("Hardware version not configured");
        return false;
    }
    std::string hardwareVersion = productConfig["hardware_version"].get<std::string>();
    if (hardwareVersion != packageInfo_.hardwareVersion) {
        LOG_ERROR("Hardware version mismatch");
        return false;
    }

    // 验证版本兼容性
    if (!productConfig.contains("version")) {
        LOG_ERROR("Product version not configured");
        return false;
    }
    std::string currentVersion = productConfig["version"].get<std::string>();
    if (currentVersion < packageInfo_.minVersion || currentVersion > packageInfo_.maxVersion) {
        LOG_ERROR("Version compatibility check failed");
        return false;
    }

    // 对于差分包，验证起始版本
    if (packageInfo_.type == UpgradeType::DIFF && currentVersion != packageInfo_.fromVersion) {
        LOG_ERROR("Current version does not match diff package from version");
        return false;
    }

    return true;
}

bool Upgrade::VerifyFileHash(const std::string &filename, const std::string &expectedHash)
{
    std::ifstream file(filename, std::ios::binary);
    if (!file) {
        LOG_ERROR("Failed to open file for hash verification: {}", filename);
        return false;
    }

    EVP_MD_CTX *ctx = EVP_MD_CTX_new();
    if (!ctx) {
        LOG_ERROR("Failed to create hash context");
        return false;
    }

    EVP_DigestInit_ex(ctx, EVP_sha256(), nullptr);

    char buf[8192];
    while (file.read(buf, sizeof(buf))) {
        EVP_DigestUpdate(ctx, buf, file.gcount());
    }
    if (file.gcount() > 0) {
        EVP_DigestUpdate(ctx, buf, file.gcount());
    }

    unsigned char hash[EVP_MAX_MD_SIZE];
    unsigned int hashLen;
    EVP_DigestFinal_ex(ctx, hash, &hashLen);
    EVP_MD_CTX_free(ctx);

    std::stringstream ss;
    for (unsigned int i = 0; i < hashLen; i++) {
        ss << std::hex << std::setw(2) << std::setfill('0') << static_cast<int>(hash[i]);
    }

    return ss.str() == expectedHash;
}

bool Upgrade::ExecuteScript(const std::string &script)
{
    std::string scriptPath = "/tmp/" + script;
    if (access(scriptPath.c_str(), F_OK) != 0) {
        LOG_ERROR("Script not found: {}", script);
        return false;
    }

    // 添加执行权限
    chmod(scriptPath.c_str(), 0755);

    // 执行脚本
    int ret = system(scriptPath.c_str());
    if (ret != 0) {
        LOG_ERROR("Script execution failed: {}", script);
        return false;
    }

    return true;
}

} // namespace Upgrade
} // namespace El